"""
文档解析服务
"""
import os
import json
import subprocess
import tempfile
import shutil
import threading
from typing import Dict, Any, Optional
from datetime import datetime
import uuid

from core.config import get_settings
from core.logger import get_task_logger

settings = get_settings()


class ParseService:
    """文档解析服务"""
    
    def __init__(self):
        self.mineru_available = self._check_mineru_availability()
        # 添加解析锁，确保同一时间只有一个解析任务在运行
        self._parse_lock = threading.Lock()
    
    def _check_mineru_availability(self) -> bool:
        """检查MinerU是否可用"""
        try:
            result = subprocess.run(
                ["python3", "-c", "import magic_pdf; print('MinerU available')"],
                capture_output=True,
                text=True,
                timeout=10
            )
            available = result.returncode == 0
            if available:
                print("MinerU is available")
            else:
                print("MinerU is not available, will use fallback methods")
            return available
        except (subprocess.TimeoutExpired, FileNotFoundError):
            print("MinerU check failed, will use fallback methods")
            return False
    
    async def parse_document(
        self,
        file_path: str,
        task_id: str,
        options: Dict[str, Any],
        task_logger
    ) -> Dict[str, Any]:
        """解析文档"""
        # 使用锁确保同一时间只有一个解析任务在运行
        task_logger.info("Acquiring parse lock", {"task_id": task_id})
        self._parse_lock.acquire()
        try:
            task_logger.info("Parse lock acquired", {"task_id": task_id})
            task_logger.info("Starting document parsing", {"file_path": file_path, "options": options})
            
            # 创建结果目录
            result_dir = os.path.join(settings.parsed_results_dir, task_id)
            os.makedirs(result_dir, exist_ok=True)
            
            results = {}
            download_urls = {}
            
            # 根据选项决定解析内容
            output_format = options.get("output_format", "all")
            
            # 至少尝试提取文本内容
            if output_format in ["all", "text"]:
                text_result = await self._extract_text(file_path, result_dir, task_logger)
                if text_result:
                    results["text"] = text_result
                    download_urls["text"] = f"/api/v1/parse/result/{task_id}?format=text"
                else:
                    task_logger.warning("Text extraction failed, trying basic extraction")
                    # 如果文本提取失败，尝试最基本的提取
                    basic_text = await self._extract_text_basic(file_path, result_dir, task_logger)
                    if basic_text:
                        results["text"] = basic_text
                        download_urls["text"] = f"/api/v1/parse/result/{task_id}?format=text"
            
            # 如果没有成功提取任何内容，至少创建一个基本的结果
            if not results:
                task_logger.warning("No content extracted successfully, creating basic result")
                basic_result = await self._create_basic_result(file_path, result_dir, task_logger)
                if basic_result:
                    results["basic"] = basic_result
                    download_urls["basic"] = f"/api/v1/parse/result/{task_id}?format=basic"
            
            # 只有在成功提取文本后才尝试其他格式
            if "text" in results:
                if output_format in ["all", "markdown"]:
                    markdown_result = await self._extract_markdown(file_path, result_dir, task_logger)
                    if markdown_result:
                        results["markdown"] = markdown_result
                        download_urls["markdown"] = f"/api/v1/parse/result/{task_id}?format=markdown"
                
                if output_format in ["all", "layout"]:
                    layout_result = await self._extract_layout(file_path, result_dir, task_logger)
                    if layout_result:
                        results["layout"] = layout_result
                        download_urls["layout"] = f"/api/v1/parse/result/{task_id}?format=layout"
                
                if output_format in ["all", "images"]:
                    images_result = await self._extract_images(file_path, result_dir, task_logger)
                    if images_result:
                        results["images"] = images_result
                        download_urls["images"] = f"/api/v1/parse/result/{task_id}?format=images"
            
            # 创建元数据文件
            metadata = {
                "task_id": task_id,
                "file_path": file_path,
                "options": options,
                "created_at": datetime.now().isoformat(),
                "results": {k: v.get("metadata", {}) for k, v in results.items()}
            }
            
            metadata_path = os.path.join(result_dir, "metadata.json")
            with open(metadata_path, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)
            
            task_logger.info("Document parsing completed", {
                "results_count": len(results),
                "result_dir": result_dir
            })
            
            return {
                "results": results,
                "download_urls": download_urls,
                "result_path": result_dir
            }
            
        except Exception as e:
            task_logger.error("Document parsing failed", error=e)
            raise
        finally:
            # 释放解析锁
            task_logger.info("Releasing parse lock", {"task_id": task_id})
            self._parse_lock.release()
            task_logger.info("Parse lock released", {"task_id": task_id})
    
    async def _extract_text(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """提取文本内容"""
        try:
            task_logger.info("Extracting text content")
            
            # 这里使用MinerU的文本提取功能
            # 由于MinerU的具体实现可能有所不同，这里提供一个通用的实现
            
            if self.mineru_available:
                # 使用MinerU提取文本
                try:
                    from magic_pdf import PDFExtractor
                    
                    extractor = PDFExtractor()
                    text_content = extractor.extract_text(file_path)
                    
                    # 保存文本结果
                    text_path = os.path.join(result_dir, "content.txt")
                    with open(text_path, 'w', encoding='utf-8') as f:
                        f.write(text_content)
                    
                    metadata = {
                        "extraction_method": "mineru",
                        "char_count": len(text_content),
                        "word_count": len(text_content.split()),
                        "extracted_at": datetime.now().isoformat()
                    }
                    
                    task_logger.info("Text extraction completed", {
                        "char_count": metadata["char_count"],
                        "word_count": metadata["word_count"]
                    })
                    
                    return {
                        "path": text_path,
                        "size": os.path.getsize(text_path),
                        "metadata": metadata
                    }
                    
                except Exception as e:
                    task_logger.warning("MinerU text extraction failed, falling back to basic method", error=e)
            
            # 备用方法：使用PyPDF2或pdfplumber
            return await self._extract_text_fallback(file_path, result_dir, task_logger)
            
        except Exception as e:
            task_logger.error("Text extraction failed", error=e)
            return None
    
    async def _extract_text_fallback(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """备用文本提取方法"""
        try:
            import pdfplumber
            
            text_content = ""
            with pdfplumber.open(file_path) as pdf:
                for page in pdf.pages:
                    text_content += page.extract_text() + "\n"
            
            # 保存文本结果
            text_path = os.path.join(result_dir, "content.txt")
            with open(text_path, 'w', encoding='utf-8') as f:
                f.write(text_content)
            
            metadata = {
                "extraction_method": "pdfplumber",
                "char_count": len(text_content),
                "word_count": len(text_content.split()),
                "extracted_at": datetime.now().isoformat()
            }
            
            task_logger.info("Text extraction completed (fallback)", {
                "char_count": metadata["char_count"],
                "word_count": metadata["word_count"]
            })
            
            return {
                "path": text_path,
                "size": os.path.getsize(text_path),
                "metadata": metadata
            }
            
        except Exception as e:
            task_logger.error("Fallback text extraction failed", error=e)
            return None
    
    async def _extract_markdown(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """提取Markdown内容"""
        try:
            task_logger.info("Extracting markdown content")
            
            if self.mineru_available:
                try:
                    from magic_pdf import PDFExtractor
                    
                    extractor = PDFExtractor()
                    markdown_content = extractor.extract_markdown(file_path)
                    
                    # 保存Markdown结果
                    markdown_path = os.path.join(result_dir, "content.md")
                    with open(markdown_path, 'w', encoding='utf-8') as f:
                        f.write(markdown_content)
                    
                    metadata = {
                        "extraction_method": "mineru",
                        "char_count": len(markdown_content),
                        "word_count": len(markdown_content.split()),
                        "extracted_at": datetime.now().isoformat()
                    }
                    
                    task_logger.info("Markdown extraction completed", {
                        "char_count": metadata["char_count"],
                        "word_count": metadata["word_count"]
                    })
                    
                    return {
                        "path": markdown_path,
                        "size": os.path.getsize(markdown_path),
                        "metadata": metadata
                    }
                    
                except Exception as e:
                    task_logger.warning("MinerU markdown extraction failed", error=e)
            
            # 备用方法：将文本转换为简单Markdown
            return await self._extract_markdown_fallback(file_path, result_dir, task_logger)
            
        except Exception as e:
            task_logger.error("Markdown extraction failed", error=e)
            return None
    
    async def _extract_markdown_fallback(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """备用Markdown提取方法"""
        try:
            # 先提取文本，然后转换为简单Markdown
            text_result = await self._extract_text_fallback(file_path, result_dir, task_logger)
            if not text_result:
                return None
            
            # 读取文本内容
            with open(text_result["path"], 'r', encoding='utf-8') as f:
                text_content = f.read()
            
            # 简单的Markdown转换
            markdown_content = self._convert_text_to_markdown(text_content)
            
            # 保存Markdown结果
            markdown_path = os.path.join(result_dir, "content.md")
            with open(markdown_path, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            
            metadata = {
                "extraction_method": "text_to_markdown",
                "char_count": len(markdown_content),
                "word_count": len(markdown_content.split()),
                "extracted_at": datetime.now().isoformat()
            }
            
            task_logger.info("Markdown extraction completed (fallback)", {
                "char_count": metadata["char_count"],
                "word_count": metadata["word_count"]
            })
            
            return {
                "path": markdown_path,
                "size": os.path.getsize(markdown_path),
                "metadata": metadata
            }
            
        except Exception as e:
            task_logger.error("Fallback markdown extraction failed", error=e)
            return None
    
    def _convert_text_to_markdown(self, text: str) -> str:
        """将文本转换为简单Markdown"""
        lines = text.split('\n')
        markdown_lines = []
        
        for line in lines:
            line = line.strip()
            if not line:
                markdown_lines.append("")
            elif line.isupper() and len(line) < 100:
                # 可能是标题
                markdown_lines.append(f"## {line}")
            else:
                markdown_lines.append(line)
        
        return '\n'.join(markdown_lines)
    
    async def _extract_layout(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """提取布局信息"""
        try:
            task_logger.info("Extracting layout information")
            
            if self.mineru_available:
                try:
                    from magic_pdf import PDFExtractor
                    
                    extractor = PDFExtractor()
                    layout_info = extractor.extract_layout(file_path)
                    
                    # 保存布局信息
                    layout_path = os.path.join(result_dir, "layout.json")
                    with open(layout_path, 'w', encoding='utf-8') as f:
                        json.dump(layout_info, f, ensure_ascii=False, indent=2)
                    
                    metadata = {
                        "extraction_method": "mineru",
                        "page_count": len(layout_info.get("pages", [])),
                        "extracted_at": datetime.now().isoformat()
                    }
                    
                    task_logger.info("Layout extraction completed", {
                        "page_count": metadata["page_count"]
                    })
                    
                    return {
                        "path": layout_path,
                        "size": os.path.getsize(layout_path),
                        "metadata": metadata
                    }
                    
                except Exception as e:
                    task_logger.warning("MinerU layout extraction failed", error=e)
            
            # 备用方法：创建简单的布局信息
            return await self._extract_layout_fallback(file_path, result_dir, task_logger)
            
        except Exception as e:
            task_logger.error("Layout extraction failed", error=e)
            return None
    
    async def _extract_layout_fallback(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """备用布局提取方法"""
        try:
            import pdfplumber
            
            layout_info = {
                "pages": []
            }
            
            with pdfplumber.open(file_path) as pdf:
                for page_num, page in enumerate(pdf.pages):
                    page_info = {
                        "page_number": page_num + 1,
                        "width": page.width,
                        "height": page.height,
                        "text_blocks": []
                    }
                    
                    # 提取文本块
                    words = page.extract_words()
                    for word in words:
                        block = {
                            "text": word["text"],
                            "x0": word["x0"],
                            "y0": word["y0"],
                            "x1": word["x1"],
                            "y1": word["y1"],
                            "width": word["x1"] - word["x0"],
                            "height": word["y1"] - word["y0"]
                        }
                        page_info["text_blocks"].append(block)
                    
                    layout_info["pages"].append(page_info)
            
            # 保存布局信息
            layout_path = os.path.join(result_dir, "layout.json")
            with open(layout_path, 'w', encoding='utf-8') as f:
                json.dump(layout_info, f, ensure_ascii=False, indent=2)
            
            metadata = {
                "extraction_method": "pdfplumber",
                "page_count": len(layout_info["pages"]),
                "extracted_at": datetime.now().isoformat()
            }
            
            task_logger.info("Layout extraction completed (fallback)", {
                "page_count": metadata["page_count"]
            })
            
            return {
                "path": layout_path,
                "size": os.path.getsize(layout_path),
                "metadata": metadata
            }
            
        except Exception as e:
            task_logger.error("Fallback layout extraction failed", error=e)
            return None
    
    async def _extract_images(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """提取图片"""
        try:
            task_logger.info("Extracting images")
            
            # 创建图片目录
            images_dir = os.path.join(result_dir, "images")
            os.makedirs(images_dir, exist_ok=True)
            
            images_info = []
            
            if self.mineru_available:
                try:
                    from magic_pdf import PDFExtractor
                    
                    extractor = PDFExtractor()
                    images = extractor.extract_images(file_path)
                    
                    for i, image_data in enumerate(images):
                        image_path = os.path.join(images_dir, f"image_{i+1}.png")
                        with open(image_path, 'wb') as f:
                            f.write(image_data["data"])
                        
                        image_info = {
                            "filename": f"image_{i+1}.png",
                            "path": image_path,
                            "size": len(image_data["data"]),
                            "page": image_data.get("page", 0),
                            "format": image_data.get("format", "png"),
                            "width": image_data.get("width", 0),
                            "height": image_data.get("height", 0)
                        }
                        images_info.append(image_info)
                    
                    task_logger.info("Image extraction completed", {
                        "image_count": len(images_info)
                    })
                    
                except Exception as e:
                    task_logger.warning("MinerU image extraction failed", error=e)
            
            # 如果没有提取到图片，尝试使用pdf2image
            if not images_info:
                images_info = await self._extract_images_fallback(file_path, images_dir, task_logger)
            
            if not images_info:
                return None
            
            # 创建图片清单
            images_list_path = os.path.join(result_dir, "images.json")
            with open(images_list_path, 'w', encoding='utf-8') as f:
                json.dump(images_info, f, ensure_ascii=False, indent=2)
            
            metadata = {
                "extraction_method": "mineru" if self.mineru_available else "pdf2image",
                "image_count": len(images_info),
                "total_size": sum(img["size"] for img in images_info),
                "extracted_at": datetime.now().isoformat()
            }
            
            return {
                "path": images_dir,
                "size": sum(img["size"] for img in images_info),
                "metadata": metadata
            }
            
        except Exception as e:
            task_logger.error("Image extraction failed", error=e)
            return None
    
    async def _extract_images_fallback(self, file_path: str, images_dir: str, task_logger) -> list:
        """备用图片提取方法"""
        try:
            from pdf2image import convert_from_path
            
            images = convert_from_path(file_path)
            images_info = []
            
            for i, image in enumerate(images):
                image_path = os.path.join(images_dir, f"page_{i+1}.png")
                image.save(image_path, 'PNG')
                
                image_info = {
                    "filename": f"page_{i+1}.png",
                    "path": image_path,
                    "size": os.path.getsize(image_path),
                    "page": i + 1,
                    "format": "png",
                    "width": image.width,
                    "height": image.height
                }
                images_info.append(image_info)
            
            task_logger.info("Image extraction completed (fallback)", {
                "image_count": len(images_info)
            })
            
            return images_info
            
        except Exception as e:
            task_logger.error("Fallback image extraction failed", error=e)
            return []
    
    async def _extract_text_basic(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """最基本的文本提取方法"""
        try:
            task_logger.info("Attempting basic text extraction")
            
            # 根据文件扩展名选择不同的处理方法
            file_ext = os.path.splitext(file_path)[1].lower()
            
            if file_ext == '.pdf':
                # 使用PyPDF2进行基本PDF文本提取
                try:
                    import PyPDF2
                    
                    text_content = ""
                    with open(file_path, 'rb') as file:
                        pdf_reader = PyPDF2.PdfReader(file)
                        for page in pdf_reader.pages:
                            text_content += page.extract_text() + "\n"
                    
                    if text_content.strip():
                        text_path = os.path.join(result_dir, "content.txt")
                        with open(text_path, 'w', encoding='utf-8') as f:
                            f.write(text_content)
                        
                        metadata = {
                            "extraction_method": "PyPDF2",
                            "char_count": len(text_content),
                            "word_count": len(text_content.split()),
                            "extracted_at": datetime.now().isoformat()
                        }
                        
                        task_logger.info("Basic text extraction completed", {
                            "method": "PyPDF2",
                            "char_count": metadata["char_count"]
                        })
                        
                        return {
                            "path": text_path,
                            "size": os.path.getsize(text_path),
                            "metadata": metadata
                        }
                except Exception as e:
                    task_logger.warning("PyPDF2 extraction failed", error=e)
            
            elif file_ext in ['.doc', '.docx']:
                # 使用python-docx提取Word文档文本
                try:
                    import docx
                    
                    doc = docx.Document(file_path)
                    text_content = ""
                    for paragraph in doc.paragraphs:
                        text_content += paragraph.text + "\n"
                    
                    if text_content.strip():
                        text_path = os.path.join(result_dir, "content.txt")
                        with open(text_path, 'w', encoding='utf-8') as f:
                            f.write(text_content)
                        
                        metadata = {
                            "extraction_method": "python-docx",
                            "char_count": len(text_content),
                            "word_count": len(text_content.split()),
                            "extracted_at": datetime.now().isoformat()
                        }
                        
                        task_logger.info("Basic text extraction completed", {
                            "method": "python-docx",
                            "char_count": metadata["char_count"]
                        })
                        
                        return {
                            "path": text_path,
                            "size": os.path.getsize(text_path),
                            "metadata": metadata
                        }
                except Exception as e:
                    task_logger.warning("python-docx extraction failed", error=e)
            
            # 如果所有方法都失败，返回None
            return None
            
        except Exception as e:
            task_logger.error("Basic text extraction failed", error=e)
            return None
    
    async def _create_basic_result(self, file_path: str, result_dir: str, task_logger) -> Optional[Dict[str, Any]]:
        """创建基本的结果文件"""
        try:
            task_logger.info("Creating basic result file")
            
            # 创建一个基本的文本文件，包含文件信息
            import os
            from datetime import datetime
            
            file_info = {
                "filename": os.path.basename(file_path),
                "file_path": file_path,
                "file_size": os.path.getsize(file_path),
                "created_at": datetime.now().isoformat(),
                "message": "Unable to extract content from this file. The file may be encrypted, corrupted, or in an unsupported format.",
                "extraction_method": "basic_info"
            }
            
            # 创建基本信息文件
            info_path = os.path.join(result_dir, "file_info.json")
            with open(info_path, 'w', encoding='utf-8') as f:
                import json
                json.dump(file_info, f, ensure_ascii=False, indent=2)
            
            # 创建一个空的文本文件
            text_path = os.path.join(result_dir, "content.txt")
            with open(text_path, 'w', encoding='utf-8') as f:
                f.write(f"# File: {file_info['filename']}\n")
                f.write(f"# Size: {file_info['file_size']} bytes\n")
                f.write(f"# Created: {file_info['created_at']}\n\n")
                f.write(file_info["message"])
            
            metadata = {
                "extraction_method": "basic_info",
                "char_count": len(text_path),
                "word_count": len(file_info["message"].split()),
                "extracted_at": datetime.now().isoformat(),
                "note": "Basic file information only - content extraction failed"
            }
            
            task_logger.info("Basic result created", {
                "filename": file_info["filename"],
                "method": "basic_info"
            })
            
            return {
                "path": text_path,
                "size": os.path.getsize(text_path),
                "metadata": metadata
            }
            
        except Exception as e:
            task_logger.error("Failed to create basic result", error=e)
            return None

parse_service = ParseService()